'''Functions to access the MySQL database.

Mike Dorey  2011-03-24
'''


import MySQLdb
import MySQLdb.cursors
import datetime


LABVIEW_EPOCH = datetime.datetime(1904, 1, 1, 0, 0, 0)


def get_connection_and_cursor(server, db_name, defaults_filename):
    '''Returns a tuple of (connection, MySQLdb.cursors.DictCursor).

    Arguments
        server - machine where the server is running
        db_name - name of the database
        defaults_filename - name of the file that holds the db user and password
    '''
    connection = MySQLdb.connect(host=server, db=db_name,
                    read_default_file=defaults_filename)
    cursor = connection.cursor(cursorclass=MySQLdb.cursors.DictCursor)
    return (connection, cursor)


def fetch_all_user_names(cursor):
    '''Returns a list of the names of all of the users in the users table.

    Argument
        cursor - a MySQLdb.cursors.DictCursor cursor
    '''
    _ = cursor.execute('select name from users')
    names = [c['name'] for c in cursor]
    return names


def fetch_user_experiment_names(cursor, user_name):
    '''Returns a list of the experiment names for the specified user.

    Arguments
        cursor - a MySQLdb.cursors.DictCursor cursor
        user_name - name of the user as it is in the users table
    '''
    _ = cursor.execute('''select e.name from experiments e, users u
                          where e.userid = u.userid and
                          u.name = %s''', (user_name,) )
    experiment_names = [c['name'] for c in cursor]
    return experiment_names


def fetch_dust_experiment(cursor, user_name, experiment_name):
    '''Returns a DustExperiment object for the specfied user and 
    experiment name.

    A a DustExperiment object holds information about the experiment 
    as well as a list of all of the dust events.

    Arguments
        cursor - a MySQLdb.cursors.DictCursor cursor
        user_name - name of the user as it is in the users table
        experiment_name - name of the experiment as it is in 
                   the experiments table
    '''
    remarks = get_remarks(cursor, user_name, experiment_name)
    dust_type = get_dust_type(cursor, experiment_name)
    dust_events = fetch_dust_events_for_experiment(cursor, user_name,
                       experiment_name)
    dust_events.sort() # ordered by timestamp
    experiment = DustExperiment(user_name, experiment_name, remarks, 
                     dust_type, dust_events)
    return experiment


def get_remarks(cursor, user_name, experiment_name):
    '''Returns the remarks for the specified experiment.

    Arguments
        cursor - a MySQLdb.cursors.DictCursor cursor
        user_name - name of the user as it is in the users table
        experiment_name - name of the experiment as it is in 
                   the experiments table
    '''
    n_rows = cursor.execute('''select e.remarks from experiments e, users u where 
                                   u.name = %s and
                                   e.name = %s and
                                   e.userid = u.userid''', 
                            (user_name, experiment_name))
    if n_rows < 1:
        return None # did not find any experiments of the specified name for the user
    c = cursor.fetchone()
    return c['remarks']


def get_dust_type(cursor, experiment_name):
    '''Returns a DustType object for the dust
    type used in the experiment.

    Arguments
        cursor - a MySQLdb.cursors.DictCursor cursor
        experiment_name - name of the experiment as it is in 
                   the experiments table
    '''
    dust_ids = get_dust_type_ids(cursor, experiment_name)
    # Just use the first one and ignore the rest. There 
    # should only be one dust type for an experiment.
    dust_id = dust_ids[0]
    dust_types = []
    __ = cursor.execute('''select chemicalcomp, batch  
                           from dusttype
                           where dustid = %s''', 
                           (dust_id,))
    c = cursor.fetchone()
    return DustType(c['chemicalcomp'], c['batch'])
   

def get_dust_type_ids(cursor, experiment_name):
    '''Returns a list of dust type ids for the specified experiment.

    Arguments
        cursor - a MySQLdb.cursors.DictCursor cursor
        experiment_name - name of the experiment as it is in 
                   the experiments table
    '''
    dust_type_times = get_dust_type_time_periods(cursor)
    exp_times = timestamp_boundaries_for_experiment(cursor, experiment_name)
    dust_type_ids = []
    for (dust_id, dust_start, dust_stop) in dust_type_times:
        for (exp_start, exp_stop) in exp_times:
            if times_overlap((dust_start, dust_stop), (exp_start, exp_stop)):
                dust_type_ids.append(dust_id)
    return list(set(dust_type_ids))


def get_dust_type_time_periods(cursor):
    '''Returns a list of (dust type id, start time, stop time) for 
    all of the dust types ever used.

    Argument
        cursor - a MySQLdb.cursors.DictCursor cursor
    '''
    n_rows = cursor.execute('''select timestamp, dusttypeid 
                               from sourcemodification 
                               order by timestamp''')
    assert n_rows > 0, 'sourcemodification table is empty'
    change_times = [(c['timestamp'], c['dusttypeid']) for c in cursor]
    return [p for p in time_periods_from_change_times(change_times)]


def times_overlap(start_stop_1, start_stop_2):
    '''Returns True if there is any overlap in the two specified
    time ranges.

    Arguments
        start_stop_1 - one time range
        start_stop_2 - a second time range
    '''
    ((a,b), (c,d)) = sorted((sorted(start_stop_1), sorted(start_stop_2)))
    return b >= c 


def current_labview_time():
    '''Returns the current time as milliseconds since the LabView epoch.'''
    now = datetime.datetime.now()
    tdiff = now - LABVIEW_EPOCH
    diff_in_seconds = tdiff.days*86400.0 + tdiff.seconds
    return long(round(diff_in_seconds*1000.0))


def time_periods_from_change_times(change_times):
    '''Generator that yields a list of (dust type id, start time, stop time) 
    derived from the times that the dust type was changed.

    Argument
        change_times - list of (timestamp, dust type id) 
    '''
    change_times = sorted(change_times)
    (t1, dtype1) = change_times[0]
    for i in range(1, len(change_times)-1):
        (t2, dtype2) = change_times[i]
        if dtype1 != dtype2:
            yield (dtype1, t1, t2)
            (t1, dtype1) = (t2, dtype2)
    # the very last change time is a special case
    (t2, dtype2) = change_times[-1]
    if dtype1 != dtype2:
        yield (dtype1, t1, t2)
        yield (dtype2, t2, current_labview_time())
    else:
        yield (dtype1, t1, current_labview_time())


def fetch_dust_events_for_experiment(cursor, user_name, experiment_name):
    '''Returns a list of DustEvent objects for the specified experiment.

    Arguments
        cursor - a MySQLdb.cursors.DictCursor cursor
        user_name - name of the user as it is in the users table
        experiment_name - name of the experiment as it is in
                   the experiments table
    '''
    user_boundaries = timestamp_boundaries_for_user(cursor, user_name)
    timestamp_strings = []
    for (start, stop) in user_boundaries:
        timestamp_strings.append('(timestamp >= %d and timestamp <= %d)' % 
                   (start, stop))
    user_where = '(' + ' or\n'.join(timestamp_strings) + ')'
    experiment_boundaries = timestamp_boundaries_for_experiment(cursor, 
                    experiment_name)
    timestamp_strings = []
    for (start, stop) in experiment_boundaries:
        timestamp_strings.append('(timestamp >= %d and timestamp <= %d)' % 
                   (start, stop))
    experiment_where = '(' + ' or\n'.join(timestamp_strings) + ')'
    where_conditions = ' and\n'.join([user_where, experiment_where])
    n_events = cursor.execute('''select * from dustevents where %s''' % 
                    where_conditions)
    return [DustEvent(c) for c in cursor]


def fetch_dust_events_that_meet_criteria(cursor, criteria):
    '''Returns a list of the dust events that meet the criteria.

    Arguments
        cursor - a MySQL.cursors.DictCursor object
        criteria - dictionary that holds the criteria
    '''
    return [DustEvent(row) for row in dust_event_gen(cursor, criteria)]


def dust_event_gen(cursor, criteria):
    '''Returns a generator that yields rows from the database that meet the 
    criteria in the dictionary. The rows are dictionaries produced by the 
    MySQL.cursors.DictCursor.

    Arguments
        cursor - a MySQL.cursors.DictCursor object
        criteria - dictionary that holds the criteria
    '''
    sql = sql_for_criteria(criteria)
    _ = cursor.execute(sql)
    for c in cursor:
        yield c


def timestamp_boundaries_for_user(cursor, user_name):
    '''Returns a list of tuples (time start, time stop) of the 
    dust events for the specified user.

    Arguments
        cursor - a MySQL.cursors.DictCursor object
        user_name - name as it appears in the users table
    '''
    n_rows = cursor.execute('''select eb.timestart, eb.timestop
                               from experimentblock eb, experiments e,
                                    users u
                               where u.name = %s and
                                     e.userid = u.userid and
                                     eb.expid = e.expid''',
                                (user_name,))
    return [(c['timestart'], c['timestop']) for c in cursor]


def timestamp_boundaries_for_experiment(cursor, experiment_name):
    '''Returns a list of tuples (time start, time stop) of the 
    dust events for the specified experiment.

    Arguments
        cursor - a MySQL.cursors.DictCursor object
        experiment_name - name as it appears in the experiments table
    '''
    n_rows = cursor.execute('''select eb.timestart, eb.timestop
                               from experimentblock eb, experiments e
                               where e.name = %s and
                                     eb.expid = e.expid''',
                                (experiment_name,))
    return [(c['timestart'], c['timestop']) for c in cursor]


def timestamp_boundaries_from_config(config):
    '''Returns a tuple (time start, time stop) as specified in the
    lower timestamp and upper timestamp values in the configuration
    dictionary.

    Argument
        config - dictionary that holds a configuration
    '''
    if config.has_key('lower timestamp'):
        if config.has_key('upper timestamp'):
            lower_and_upper_timestamps = (long(config['lower timestamp']), 
                                          long(config['upper timestamp']))
        else:
            # has a lower timestamp but does not have an upper timestamp
            lower_and_upper_timestamps = (long(config['lower timestamp']),
                                          current_labview_time())
    else: # has no lower timestamp
        if config.has_key('upper timestamp'):
            # has no lower timestamp but does have an upper timestamp
            lower_and_upper_timestamps = (0, long(config['upper timestamp']))
        else:
            # has neither a lower or upper timestamp
            lower_and_upper_timestamps = None
    return lower_and_upper_timestamps


def timestamp_boundaries_for_dust_type(cursor, dust_type):
    '''Returns a list of tuples (time start, time stop) of the 
    dust events for specified dust type.

    Arguments
        cursor - a MySQL.cursors.DictCursor object
        dust_type - DustType object
    '''
    id_start_stops = get_dust_type_time_periods(cursor)
    n_rows = cursor.execute('''select dustid from dusttype 
                               where chemicalcomp = %s and
                                     batch = %s''',
                               (dust_type.chemical_comp, dust_type.batch))
    dust_ids = [c['dustid'] for c in cursor]
    starts_stops = []
    for dust_id in dust_ids:
        for the_id, the_start, the_stop in id_start_stops:
            if the_id == dust_id:
                starts_stops.append((the_start, the_stop))
    return starts_stops


def sql_for_criteria(criteria):
    '''Returns the SQL statement that selects the dust events for
    a single experiment for a particular user. These dust events 
    are selected so that they meet the criteria in the dictionary.

    Argument
        criteria - dictionary that holds the criteria
    '''
    where_parts = []
    if criteria.has_key('timestamp_groups'):
        groups_strings = []
        for group in criteria['timestamp_groups']:
            group_strings = []
            for (start, stop) in group:
                s = '(timestamp >= %d and timestamp < %d)' % \
                      (start, stop)
                group_strings.append(s)
            group_string = '(' + ' or\n'.join(group_strings) + ')'
            groups_strings.append(group_string)
        timestamp_where = ' and\n'.join(groups_strings)
        where_parts.append(timestamp_where)
    if criteria.has_key('lower_and_upper_velocities'):
        velocities_where = ' or\n'.join(['(velocity >= %s and velocity < %s)' %
                                  (l, u)  for (l, u) in 
                                  criteria.get('lower_and_upper_velocities', 
                                               [])])
        where_parts.append(velocities_where)
    if criteria.has_key('lower_and_upper_charges'):
        charges_where = ' or\n'.join(['(charge >= %s and charge < %s)' %
                                  (l, u)  for (l, u) in 
                                  criteria.get('lower_and_upper_charges', [])])
        where_parts.append(charges_where)
    if criteria.has_key('lower_and_upper_masses'):
        masses_where = ' or\n'.join(['(mass >= %s and mass < %s)' %
                                  (l, u)  for (l, u) in 
                                  criteria.get('lower_and_upper_masses', [])])
        where_parts.append(masses_where)
    if criteria.has_key('selection'):
        selection_where = 'selection = ' + str(criteria['selection'])
        where_parts.append(selection_where)
    where_part = ' and\n'.join(where_parts)
    if len(where_part) > 0:
        where_section = 'where ' + where_part
    else:
        where_section = ''
    sql = '''select timestamp, velocity, charge,
                    mass, terminalpotential,
                    lenspotential, needlepotential,
                    pulsefreq, pulsewidth,
                    pulseheight, selection,
                    waveforms
             from dustevents 
             ''' + \
             where_section
    return sql


class DustExperiment(object):
    '''Objects of this class are dust experiments.'''

    def __init__(self, user_name, experiment_name, remarks, dust_type, 
             dust_events):
        '''Construct a dust experiment from its parts.

        Arguments
            user_name - name of the user
            experiment_name - name of the experiment
            remarks - remarks about the experiment
            timestart - when the experiment started
            timestop - when the experiment stopped
            dust_type - the DustType used in this experiment
            dust_events - list of dust events for the experiment
        '''
        self.user_name = user_name
        self.experiment_name = experiment_name
        self.remarks = remarks
        self.dust_type = dust_type
        self.dust_events = dust_events


    def __str__(self):
        '''Returns a string representation of this experiment.'''
        return 'user name: %s\texperiment name: %s\nchemical comp: %s\tbatch: %s' % \
               (self.user_name, self.experiment_name, 
                self.dust_type.chemical_comp, self.dust_type.batch)


    __repr__ = __str__


class DustEvent(object):
    '''Objects of this class are dust events fetched from the database.'''


    def __init__(self, dct):
        '''Constructs a DustEvent object using key/values pairs in the 
        dictionary.

        Argument
            dct - dictionary that is probably returned by a DictCursor.
                Its keys are the fields of the dustevents table.
        '''
        self.timestamp = dct['timestamp']
        self.velocity = dct['velocity']
        self.charge = dct['charge']
        self.mass = dct['mass']
        self.terminalpotential = dct['terminalpotential']
        self.lenspotential = dct['lenspotential']
        self.needlepotential = dct['needlepotential']
        self.pulsefreq = dct['pulsefreq']
        self.pulsewidth = dct['pulsewidth']
        self.pulseheight = dct['pulseheight']
        self.selection = dct['selection']
        self.waveforms = dct['waveforms']

   
    def __cmp__(self, other):
        '''Compares two dust events based on their timestamps.'''
        return cmp(self.timestamp, other.timestamp)


    def __str__(self):
        '''Returns a string representation of this object.'''
        return '%d %f %f %f' % (self.timestamp, self.velocity,
               self.charge, self.mass)


    __repr__ = __str__


class DustType(object):
    '''Objects of this class are types of dust.'''


    def __init__(self, chemical_comp, batch):
        '''Constructs a dust type from its parts.

        Arguments
            chemical_comp - something like 'IRON' or 'NICKEL'
            batch - something like '64' or '112'
        '''
        self.chemical_comp = chemical_comp
        self.batch = batch


    def __str__(self):
        '''Returns a string representation of this object.'''
        return '%s %s' % (self.chemical_comp,
             self.batch)


    __repr__ = __str__
